import json
import requests
import websocket
import streamlit as st

ACCESS_TOKEN = 'sk-4YmeWDNiRWkbl1QF7WdcEmZkmdAAtbqcx0qMnBbnA5QAzuQ2'
BASE_URL = 'https://api.tgkwai.com'


def initialize_session_state():
    keys = ["session_id", "messages", "uploaded_file", "attachment_url"]
    for key in keys:
        if key not in st.session_state:
            st.session_state[key] = None if key != "messages" else []


def upload_attachment(file):
    try:
        response = requests.post(
            f"{BASE_URL}/api/v1/attachments/upload",
            headers={'X-Token': f'Bearer {ACCESS_TOKEN}'},
            files={'type': (None, 'image'), 'file': file}
        )
        return response.json()['data'] if response.json()['code'] == 200 else None
    except Exception as e:
        st.error(f"上传失败: {str(e)}")
        return None


def delete_attachment(file_path):
    try:
        requests.delete(
            f"{BASE_URL}/api/v1/attachments/delete",
            headers={'X-Token': f'Bearer {ACCESS_TOKEN}'},
            params={'file_path': file_path}
        )
    except Exception as e:
        st.error(f"删除失败: {str(e)}")


def create_websocket(session_id=None):
    try:
        ws_url = f"wss://api.tgkwai.com/api/v1/qamodel/session?x-token={ACCESS_TOKEN}"
        if session_id:
            ws_url += f"&session_id={session_id}"
        ws = websocket.create_connection(ws_url)
        response = json.loads(ws.recv())
        if response.get('code') != 200:
            st.error(f"连接失败: {response.get('message', '未知错误')}")
            return None, None
        return ws, response['data']['session_id']  # 确保正确获取session_id
    except Exception as e:
        st.error(f"连接异常: {str(e)}")
        return None, None


def handle_streaming_response(ws):
    full_content = []
    try:
        while True:
            result = ws.recv()
            data = json.loads(result)
            if data.get('code') != 200:
                st.error(f"API错误: {data.get('message', '未知错误')}")
                break

            content = data['data'].get('content', '')
            if content:
                full_content.append(content)

            if data['data'].get('is_finished', False):
                break
    except websocket.WebSocketConnectionClosedException:
        st.error("连接意外关闭")
    except Exception as e:
        st.error(f"接收异常: {str(e)}")
    return ''.join(full_content)


def main():
    st.set_page_config("作物病害诊断系统", "🌾", layout="centered")
    initialize_session_state()

    st.title("🌱 智能作物诊断系统")

    # 文件上传模块
    with st.sidebar:
        uploaded_file = st.file_uploader("上传作物图片", type=["jpg", "jpeg", "png"])
        if uploaded_file and uploaded_file != st.session_state.uploaded_file:
            if st.session_state.attachment_url:
                delete_attachment(st.session_state.attachment_url)
            with st.spinner("上传中..."):
                st.session_state.attachment_url = upload_attachment(uploaded_file)
                st.session_state.uploaded_file = uploaded_file
            st.rerun()

    # 聊天显示模块
    for msg in st.session_state.messages:
        with st.chat_message(msg["role"]):
            if msg.get("type") == "multimodal":
                st.image(st.session_state.uploaded_file, width=200)
                st.markdown(msg["content"])
            else:
                st.markdown(msg["content"])

    if prompt := st.chat_input("输入您的问题..."):
        new_message = {"role": "user", "content": prompt}

        # 首次提问时附加图片
        if st.session_state.attachment_url and not st.session_state.messages:
            new_message.update({
                "type": "multimodal",
                "attachments": [{"type": "image", "url": st.session_state.attachment_url}]
            })

        # 更新消息历史
        st.session_state.messages.append(new_message)
        with st.chat_message("user"):
            if new_message.get("type") == "multimodal":
                st.image(st.session_state.uploaded_file, width=200)
            st.markdown(prompt)

        # 显示等待中的助手消息
        with st.chat_message("assistant"):
            placeholder = st.empty()
            placeholder.markdown("▌")

        # 发送请求
        with st.spinner("分析中..."):
            ws, session_id = create_websocket(st.session_state.session_id)
            if ws:
                try:
                    # 确保发送正确的消息格式
                    request_data = {
                        "session_id": session_id,
                        "messages": [msg for msg in st.session_state.messages if msg["role"] != "assistant"],
                        "stream": True
                    }
                    ws.send(json.dumps(request_data))

                    # 处理流式响应
                    full_response = []
                    while True:
                        result = ws.recv()
                        data = json.loads(result)
                        if data.get('code') != 200:
                            st.error(f"API错误: {data.get('message', '未知错误')}")
                            break

                        chunk = data['data'].get('content', '')
                        if chunk:
                            full_response.append(chunk)
                            placeholder.markdown(''.join(full_response) + "▌")

                        if data['data'].get('is_finished', False):
                            break

                    # 更新最终响应
                    final_response = ''.join(full_response)
                    placeholder.markdown(final_response)
                    st.session_state.messages[-1]["content"] = final_response
                    st.session_state.session_id = session_id

                except Exception as e:
                    st.error(f"请求处理异常: {str(e)}")
                finally:
                    ws.close()


if __name__ == "__main__":
    main()
